package com.woniu.service;

import com.woniu.utils.ConversationContextManager;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.converter.BeanOutputConverter;
import org.springframework.ai.tool.ToolCallbackProvider;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;


/**
 * @Author: 朱杰
 * @Description: 自定义的服务类，用于处理用户输入并返回响应
 * @DateTime: 2025/6/10 10:35
 **/
@Service
public class ChatService {
    private final ChatClient chatClient;
    @Autowired
    private ConversationContextManager contextManager;

    public ChatService(ChatClient.Builder chatClientBuilder, ToolCallbackProvider tools) {
        this.chatClient = chatClientBuilder
                .defaultToolCallbacks(tools)
                .defaultSystem("""
                    你是一个专业的网络安全靶场考试AI助手，主要职责包括：
                    1. 解答学生在CTF考试、渗透测试、漏洞分析中遇到的问题
                    2. 如果学生需要便帮助学生创建/删除训练靶场环境（需调用dockerService工具）
                    3. 执行安全的数据库操作（需调用databaseService工具，调用该工具的方法时需要进行权限验证，我会提供当前用户的userId给你，不要提醒用户让他确认他的用户id）
                    4. 分析学生考试表现，提供针对性学习建议
                    5. 时间戳按yyyy-MM-dd HH:mm:ss转换
                    6. 在返回结果时，不要展示敏感信息，如密码、创建时间、创建人、更新时间、更新人、删除时间、删除人等
                    
                    注意事项：
                    -对返回结果进行格式化，不要返回原始数据
                    - 如果从数据库中查询的数据中有时间戳，需要将其格式化为“yyyy-MM-dd HH:mm:ss”的格式
                    - 不要返回任何敏感信息，如密码、密钥等
                    - 不要返回任何可能导致系统崩溃或安全漏洞的信息
                    """)
                .defaultTools()
                .build();
    }

    public Flux<String> processQuery(String userId, String userInput) {
        BeanOutputConverter<String> parser = new BeanOutputConverter<>(String.class);
        contextManager.addUserInput(userId, userInput);
        List<String> context = contextManager.getContext(userId);
        String fullInput = String.join("\n", context);
        fullInput = fullInput + "userId=" + userId;
        Flux<String> responseStream = chatClient.prompt(fullInput).user(u -> u.text(userInput)
                .param("format", parser.getFormat())
                .param("structure", "使用自然段落格式，每个段落之间用两个换行符分隔"))
                .stream()
                .content()
                .map(content ->
                     content.replace("\n\n", "<br/><br/>")
                             .replace("\n", "<br/>"));

        StringBuilder aiResponseBuilder = new StringBuilder();
        return responseStream.doOnNext(aiResponseBuilder::append)
                .doOnComplete(() -> {
                    String aiResponse = aiResponseBuilder.toString();
                    contextManager.addAIResponse(userId, aiResponse);
                });
    }

    public void clearContext(String userId) {
        contextManager.clearContext(userId);
    }
}


